{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "# Transformer的PyTorch实现\n",
    "\n",
    "> 本文由**罗周杨 stupidme.me.lzy@gmail.com**原创，转载请注明原作者和出处。\n",
    "> 本文严禁用于商业用途。\n",
    "\n",
    "之前研究了一番**Transformer**模型，也有一个笔记[Transformer学习笔记](transformer.ipynb)，现在我们用PyTorch实现。\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Word embedding的实现\n",
    "\n",
    "这个已经是seq2seq最基本的操作了。PyTorch也有实现，直接调用即可：\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "\n",
    "\n",
    "embedding = nn.Embedding(vocab_size, embedding_size, padding_idx=0)\n",
    "seq_embedding = seq_embedding(inputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Position encoding的实现\n",
    "再回顾一下什么是Position encoding。我们知道Transformer模型由于没有使用RNN，必须使用额外的手段来获取文本序列的顺序（或者说位置）信息。我们的Position encoding(PE)就是实现这个目标而出现的。\n",
    "\n",
    "Position encoding和Word embedding十分相似。Word embeddings是对词语的内容进行嵌入，而Position encoding是对词语的位置进行嵌入。所以，Position encoding还有个很流行的名字，叫做Position embedding。\n",
    "\n",
    "下面我们看代码吧，然后讲解代码，你就可以弄懂其中的所有细节了。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "\n",
    "class PositionalEncoding(nn.Module):\n",
    "    \n",
    "    def __init__(self, d_model, max_seq_len):\n",
    "        \"\"\"初始化。\n",
    "        \n",
    "        Args:\n",
    "            d_model: 一个标量。模型的维度，论文默认是512\n",
    "            max_seq_len: 一个标量。文本序列的最大长度\n",
    "        \"\"\"\n",
    "        super(PositionalEncoding, self).__init__()\n",
    "        \n",
    "        # 根据论文给的公式，构造出PE矩阵\n",
    "        position_encoding = np.array([\n",
    "          [pos / np.pow(10000, 2.0 * (j // 2) / d_model) for j in range(d_model)]\n",
    "          for pos in range(max_seq_len)])\n",
    "        # 偶数列使用sin，奇数列使用cos\n",
    "        position_encoding[:, 0::2] = np.sin(position_encoding[:, 0::2])\n",
    "        position_encoding[:, 1::2] = np.cos(position_encoding[:, 1::2])\n",
    "\n",
    "        # 在PE矩阵的第一行，加上一行全是0的向量，代表这`PAD`的positional encoding\n",
    "        # 在word embedding中也经常会加上`UNK`，代表位置单词的word embedding，两者十分类似\n",
    "        # 那么为什么需要这个额外的PAD的编码呢？很简单，因为文本序列的长度不一，我们需要对齐，\n",
    "        # 短的序列我们使用0在结尾补全，我们也需要这些补全位置的编码，也就是`PAD`对应的位置编码\n",
    "        pad_row = torch.zeros([1, d_model])\n",
    "        position_encoding = torch.cat((pad_row, position_encoding))\n",
    "        \n",
    "        # 嵌入操作，+1是因为增加了`PAD`这个补全位置的编码，\n",
    "        # Word embedding中如果词典增加`UNK`，我们也需要+1。看吧，两者十分相似\n",
    "        self.position_encoding = nn.Embedding(max_seq_len + 1, d_model)\n",
    "        self.position_encoding.weight = nn.Parameter(position_encoding,\n",
    "                                                     requires_grad=False)\n",
    "    def forward(self, input_len):\n",
    "        \"\"\"神经网络的前向传播。\n",
    "\n",
    "        Args:\n",
    "          input_len: 一个张量，形状为[BATCH_SIZE, 1]。每一个张量的值代表这一批文本序列中对应的长度。\n",
    "\n",
    "        Returns:\n",
    "          返回这一批序列的位置编码，进行了对齐。\n",
    "        \"\"\"\n",
    "        \n",
    "        # 找出这一批序列的最大长度\n",
    "        max_len = torch.max(input_len)\n",
    "        tensor = torch.cuda.LongTensor if input_len.is_cuda else torch.LongTensor\n",
    "        # 对每一个序列的位置进行对齐，在原序列位置的后面补上0\n",
    "        # 这里range从1开始也是因为要避开PAD(0)的位置\n",
    "        input_pos = tensor(\n",
    "          [list(range(1, len + 1)) + [0] * (max_len - len) for len in input_len])\n",
    "        return self.position_encoding(input_pos)\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Layer normalization的实现\n",
    "Layer nomalization和Batch normalization的比较，在文章开头提到的笔记已经有说明了。\n",
    "\n",
    "接下来就来实现它吧。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "\n",
    "class LayerNorm(nn.Module):\n",
    "    \"\"\"实现LayerNorm。其实PyTorch已经实现啦，见nn.LayerNorm。\"\"\"\n",
    "\n",
    "    def __init__(self, features, epsilon=1e-6):\n",
    "        \"\"\"Init.\n",
    "\n",
    "        Args:\n",
    "            features: 就是模型的维度。论文默认512\n",
    "            epsilon: 一个很小的数，防止数值计算的除0错误\n",
    "        \"\"\"\n",
    "        super(LayerNorm, self).__init__()\n",
    "        # weights\n",
    "        self.gamma = nn.Parameter(torch.ones(features))\n",
    "        # bias\n",
    "        self.beta = nn.Parameter(torch.zeros(features))\n",
    "        self.epsilon = epsilon\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"前向传播.\n",
    "\n",
    "        Args:\n",
    "            x: 输入序列张量，形状为[B, L, D]\n",
    "        \"\"\"\n",
    "        # 根据公式进行归一化\n",
    "        # 在X的最后一个维度求均值，最后一个维度就是模型的维度\n",
    "        mean = x.mean(-1, keepdim=True)\n",
    "        # 在X的最后一个维度求方差，最后一个维度就是模型的维度\n",
    "        std = x.std(-1, keepdim=True)\n",
    "        return self.gamma * (x - mean) / (std + self.epsilon) + self.beta\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Scaled dot-product attention的实现\n",
    "\n",
    "具体的原理在另一个笔记有提到。咱们现在实现它，这个代码也比较简单，根据论文的公式即可。唯一要注意的是，有个mask。\n",
    "\n",
    "代码如下："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "\n",
    "class ScaledDotProductAttention(nn.Module):\n",
    "    \"\"\"Scaled dot-product attention mechanism.\"\"\"\n",
    "\n",
    "    def __init__(self, attention_dropout=0.0):\n",
    "        super(ScaledDotProductAttention, self).__init__()\n",
    "        self.dropout = nn.Dropout(attention_dropout)\n",
    "        self.softmax = nn.Softmax(dim=2)\n",
    "\n",
    "    def forward(self, q, k, v, scale=None, attn_mask=None):\n",
    "        \"\"\"前向传播.\n",
    "\n",
    "        Args:\n",
    "          q: Queries张量，形状为[B, L_q, D_q]\n",
    "          k: Keys张量，形状为[B, L_k, D_k]\n",
    "          v: Values张量，形状为[B, L_v, D_v]，一般来说就是k\n",
    "          scale: 缩放因子，一个浮点标量\n",
    "          attn_mask: Masking张量，形状为[B, L_q, L_k]\n",
    "\n",
    "        Returns:\n",
    "          上下文张量和attetention张量\n",
    "        \"\"\"\n",
    "        attention = torch.bmm(q, k.transpose(1, 2))\n",
    "        if scale:\n",
    "          attention = attention * scale\n",
    "        if attn_mask:\n",
    "          # 给需要mask的地方设置一个负无穷\n",
    "          attention = attention.masked_fill_(attn_mask, -np.inf)\n",
    "        attention = self.softmax(attention)\n",
    "        attention = self.dropout(attention)\n",
    "        context = torch.bmm(attention, v)\n",
    "        return context, attention\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Multihead attention的实现\n",
    "\n",
    "重点戏来啦，那就是multihead attention的实现。分析可以看文章开始提到的笔记。\n",
    "\n",
    "代码如下：\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "\n",
    "class MultiHeadAttention(nn.Module):\n",
    "\n",
    "    def __init__(self, model_dim=512, num_heads=8, dropout=0.0):\n",
    "        super(MultiHeadAttention, self).__init__()\n",
    "\n",
    "        self.dim_per_head = model_dim // num_heads\n",
    "        self.num_heads = num_heads\n",
    "        self.linear_k = nn.Linear(model_dim, self.dim_per_head * num_heads)\n",
    "        self.linear_v = nn.Linear(model_dim, self.dim_per_head * num_heads)\n",
    "        self.linear_q = nn.Linear(model_dim, self.dim_per_head * num_heads)\n",
    "\n",
    "        self.dot_product_attention = ScaledDotProductAttention(dropout)\n",
    "        self.linear_final = nn.Linear(model_dim, model_dim)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        self.layer_norm = nn.LayerNorm(model_dim)\n",
    "\n",
    "    def forward(self, key, value, query, attn_mask=None):\n",
    "\n",
    "        residual = query\n",
    "\n",
    "        dim_per_head = self.dim_per_head\n",
    "        num_heads = self.num_heads\n",
    "        batch_size = key.size(0)\n",
    "\n",
    "        # linear projection\n",
    "        key = self.linear_k(key)\n",
    "        value = self.linear_v(value)\n",
    "        query = self.linear_q(query)\n",
    "\n",
    "        # split by heads\n",
    "        key = key.view(batch_size * num_heads, -1, dim_per_head)\n",
    "        value = value.view(batch_size * num_heads, -1, dim_per_head)\n",
    "        query = query.view(batch_size * num_heads, -1, dim_per_head)\n",
    "\n",
    "        if attn_mask:\n",
    "            attn_mask = attn_mask.repeat(num_heads, 1, 1)\n",
    "        # scaled dot product attention\n",
    "        scale = (key.size(-1) // num_heads) ** -0.5\n",
    "        context, attention = self.dot_product_attention(\n",
    "          query, key, value, scale, attn_mask)\n",
    "\n",
    "        # concat heads\n",
    "        context = context.view(batch_size, -1, dim_per_head * num_heads)\n",
    "\n",
    "        # final linear projection\n",
    "        output = self.linear_final(context)\n",
    "\n",
    "        # dropout\n",
    "        output = self.dropout(output)\n",
    "\n",
    "        # add residual and norm layer\n",
    "        output = self.layer_norm(residual + output)\n",
    "\n",
    "        return output, attention\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Positional-wise feed forward network的实现\n",
    "\n",
    "实际上就是一维卷积。代码如下："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "\n",
    "class PositionalWiseFeedForward(nn.Module):\n",
    "\n",
    "    def __init__(self, model_dim=512, ffn_dim=2048, dropout=0.0):\n",
    "        super(PositionalWiseFeedForward, self).__init__()\n",
    "        self.w1 = nn.Conv1d(model_dim, ffn_dim, 1)\n",
    "        self.w2 = nn.Conv1d(model_dim, ffn_dim, 1)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        self.layer_norm = nn.LayerNorm(model_dim)\n",
    "\n",
    "    def forward(self, x):\n",
    "        output = x.transpose(1, 2)\n",
    "        output = self.w2(F.relu(self.w1(output)))\n",
    "        output = self.dropout(output.transpose(1, 2))\n",
    "\n",
    "        # add residual and norm layer\n",
    "        output = self.layer_norm(x + output)\n",
    "        return output\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Encoder的实现\n",
    "\n",
    "编码器有6层，每一层的结构为EncoderLayer，总的编码器为Encoder。\n",
    "\n",
    "代码如下："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "\n",
    "class EncoderLayer(nn.Module):\n",
    "\n",
    "    def __init__(self, model_dim=512, num_heads=8, ffn_dim=2018, dropout=0.0):\n",
    "        super(EncoderLayer, self).__init__()\n",
    "\n",
    "        self.attention = MultiHeadAttention(model_dim, num_heads, dropout)\n",
    "        self.feed_forward = PositionalWiseFeedForward(model_dim, ffn_dim, dropout)\n",
    "\n",
    "    def forward(self, inputs, attn_mask=None):\n",
    "\n",
    "        # self attention\n",
    "        context, attention = self.attention(inputs, inputs, inputs, padding_mask)\n",
    "\n",
    "        # feed forward network\n",
    "        output = self.feed_forward(context)\n",
    "\n",
    "        return output, attention\n",
    "\n",
    "\n",
    "class Encoder(nn.Module):\n",
    "\n",
    "    def __init__(self,\n",
    "               vocab_size,\n",
    "               max_seq_len,\n",
    "               num_layers=6,\n",
    "               model_dim=512,\n",
    "               num_heads=8,\n",
    "               ffn_dim=2048,\n",
    "               dropout=0.0):\n",
    "        super(Encoder, self).__init__()\n",
    "\n",
    "        self.encoder_layers = nn.ModuleList(\n",
    "          [EncoderLayer(model_dim, num_heads, ffn_dim, dropout) for _ in\n",
    "           range(num_layers)])\n",
    "\n",
    "        self.seq_embedding = nn.Embedding(vocab_size + 1, model_dim, padding_idx=0)\n",
    "        self.pos_embedding = PositionalEncoding(model_dim, max_seq_len)\n",
    "\n",
    "    def forward(self, inputs, inputs_len):\n",
    "        output = self.seq_embedding(inputs)\n",
    "        output += self.pos_embedding(inputs_len)\n",
    "\n",
    "        self_attention_mask = padding_mask(inputs, inputs)\n",
    "\n",
    "        attentions = []\n",
    "        for encoder in self.encoder_layers:\n",
    "            output, attention = encoder(output, self_attention_mask)\n",
    "            attentions.append(attention)\n",
    "\n",
    "        return output, attentions\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Decoder的实现\n",
    "\n",
    "Decoder也有6层，每一层的结构为DecoderLayer，整个解码器为Decoder。\n",
    "\n",
    "代码如下："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "\n",
    "class DecoderLayer(nn.Module):\n",
    "\n",
    "    def __init__(self, model_dim, num_heads=8, ffn_dim=2048, dropout=0.0):\n",
    "        super(DecoderLayer, self).__init__()\n",
    "\n",
    "        self.attention = MultiHeadAttention(model_dim, num_heads, dropout)\n",
    "        self.feed_forward = PositionalWiseFeedForward(model_dim, ffn_dim, dropout)\n",
    "\n",
    "    def forward(self,\n",
    "              dec_inputs,\n",
    "              enc_outputs,\n",
    "              self_attn_mask=None,\n",
    "              context_attn_mask=None):\n",
    "        # self attention, all inputs are decoder inputs\n",
    "        dec_output, self_attention = self.attention(\n",
    "          dec_inputs, dec_inputs, dec_inputs, self_attn_mask)\n",
    "\n",
    "        # context attention\n",
    "        # query is decoder's outputs, key and value are encoder's inputs\n",
    "        dec_output, context_attention = self.attention(\n",
    "          enc_outputs, enc_outputs, dec_output, context_attn_mask)\n",
    "\n",
    "        # decoder's output, or context\n",
    "        dec_output = self.feed_forward(dec_output)\n",
    "\n",
    "        return dec_output, self_attention, context_attention\n",
    "\n",
    "\n",
    "class Decoder(nn.Module):\n",
    "\n",
    "    def __init__(self,\n",
    "               vocab_size,\n",
    "               max_seq_len,\n",
    "               num_layers=6,\n",
    "               model_dim=512,\n",
    "               num_heads=8,\n",
    "               ffn_dim=2048,\n",
    "               dropout=0.0):\n",
    "        super(Decoder, self).__init__()\n",
    "\n",
    "        self.num_layers = num_layers\n",
    "\n",
    "        self.decoder_layers = nn.ModuleList(\n",
    "          [DecoderLayer(model_dim, num_heads, ffn_dim, dropout) for _ in\n",
    "           range(num_layers)])\n",
    "\n",
    "        self.seq_embedding = nn.Embedding(vocab_size + 1, model_dim, padding_idx=0)\n",
    "        self.pos_embedding = PositionalEncoding(model_dim, max_seq_len)\n",
    "\n",
    "    def forward(self, inputs, inputs_len, enc_output, context_attn_mask=None):\n",
    "        output = self.seq_embedding(inputs)\n",
    "        output += self.pos_embedding(inputs_len)\n",
    "\n",
    "        self_attention_padding_mask = padding_mask(inputs, inputs)\n",
    "        seq_mask = sequence_mask(inputs)\n",
    "        self_attn_mask = torch.gt((self_attention_padding_mask + seq_mask), 0)\n",
    "\n",
    "        self_attentions = []\n",
    "        context_attentions = []\n",
    "        for decoder in self.decoder_layers:\n",
    "            output, self_attn, context_attn = decoder(\n",
    "            output, enc_output, self_attn_mask, context_attn_mask)\n",
    "            self_attentions.append(self_attn)\n",
    "            context_attentions.append(context_attn)\n",
    "\n",
    "        return output, self_attentions, context_attentions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Padding mask和sequence mask的实现\n",
    "\n",
    "笔记有说到，这里的mask有两种。\n",
    "\n",
    "代码如下："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "\n",
    "def sequence_mask(seq):\n",
    "    batch_size, seq_len = seq.size()\n",
    "    mask = torch.triu(torch.ones((seq_len, seq_len), dtype=torch.uint8),\n",
    "                    diagonal=1)\n",
    "    mask = mask.unsqueeze(0).expand(batch_size, -1, -1)  # [B, L, L]\n",
    "    return mask\n",
    "\n",
    "\n",
    "def padding_mask(seq_k, seq_q):\n",
    "    len_q = seq_q.size(1)\n",
    "    # `PAD` is 0\n",
    "    pad_mask = seq_k.eq(0)\n",
    "    pad_mask = pad_mask.unsqueeze(1).expand(-1, len_q, -1)  # shape [B, L_q, L_k]\n",
    "    return pad_mask\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Transformer的实现\n",
    "\n",
    "现在把encoder和decoder组合起来，就行啦。\n",
    "\n",
    "代码如下："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "\n",
    "class Transformer(nn.Module):\n",
    "\n",
    "    def __init__(self,\n",
    "               src_vocab_size,\n",
    "               src_max_len,\n",
    "               tgt_vocab_size,\n",
    "               tgt_max_len,\n",
    "               num_layers=6,\n",
    "               model_dim=512,\n",
    "               num_heads=8,\n",
    "               ffn_dim=2048,\n",
    "               dropout=0.2):\n",
    "        super(Transformer, self).__init__()\n",
    "\n",
    "        self.encoder = Encoder(src_vocab_size, src_max_len, num_layers, model_dim,\n",
    "                               num_heads, ffn_dim, dropout)\n",
    "        self.decoder = Decoder(tgt_vocab_size, tgt_max_len, num_layers, model_dim,\n",
    "                               num_heads, ffn_dim, dropout)\n",
    "\n",
    "        self.linear = nn.Linear(model_dim, tgt_vocab_size, bias=False)\n",
    "        self.softmax = nn.Softmax(dim=2)\n",
    "\n",
    "    def forward(self, src_seq, src_len, tgt_seq, tgt_len):\n",
    "        context_attn_mask = padding_mask(tgt_seq, src_seq)\n",
    "\n",
    "        output, enc_self_attn = self.encoder(src_seq, src_len)\n",
    "\n",
    "        output, dec_self_attn, ctx_attn = self.decoder(\n",
    "          tgt_seq, tgt_len, output, context_attn_mask)\n",
    "\n",
    "        output = self.linear(output)\n",
    "        output = self.softmax(output)\n",
    "\n",
    "        return output, enc_self_attn, dec_self_attn, ctx_attn\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "花了好几天时间读文章和看代码，基本上搞懂了，但是到了这里就没有写出来的欲望了——因为感觉好累。。。\n",
    "\n",
    "完整的代码：[transformer_pytorch.py](https://github.com/luozhouyang/machine-learning-notes/blob/master/transformer/transformer_pytorch.py)\n",
    "\n",
    "贴一些有用的链接供大家参考：\n",
    "\n",
    "* [Building the Mighty Transformer for Sequence Tagging in PyTorch : Part I](https://medium.com/@kolloldas/building-the-mighty-transformer-for-sequence-tagging-in-pytorch-part-i-a1815655cd8)\n",
    "* [Building the Mighty Transformer for Sequence Tagging in PyTorch : Part II](https://medium.com/@kolloldas/building-the-mighty-transformer-for-sequence-tagging-in-pytorch-part-ii-c85bf8fd145)\n",
    "* [The Annotated Transformer](http://nlp.seas.harvard.edu/2018/04/03/attention.html)\n",
    "* [jadore801120/attention-is-all-you-need-pytorch](https://github.com/jadore801120/attention-is-all-you-need-pytorch)\n",
    "* [JayParks/transformer](https://github.com/JayParks/transformer)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
